import sys
import codecs
import os
import matplotlib.pyplot as plt
import numpy as np

log_dir = sys.argv[1]
log_files = [fname for fname in os.listdir(log_dir) if ".log" in fname]

for filename in log_files:
    accs, precs, recs, f1s = [], [], [], []
    log_file = os.path.join(log_dir, filename)
    print("log file:", log_file)
    with codecs.open(log_file, 'r', encoding= u'utf-8',errors='ignore') as fr:
        for line in fr:
            if "Post-MetaTrain Performance of model" in line:
                s = line.split(":")[-1].strip().strip("(").strip(")").split(",", 2)
                accs.append(float(s[0]))
                P_arr, R_arr, F_arr, cnt_arr = s[-1].strip("(").strip(")").split("]),")
                precs.append(np.mean([float(item) for item in P_arr.strip().strip("array([").split(",")]))
                recs.append(np.mean([float(item) for item in R_arr.strip().strip("array([").split(",")]))
                f1s.append(np.mean([float(item) for item in F_arr.strip().strip("array([").split(",")]))
                pos_cnt, nature_cnt, neg_cnt = [int(item) for item in cnt_arr.strip().strip("]").strip("array([").split(",")]


    steps = list(range(len(accs)))
    plt.cla()
    plt.title(f"{filename.split('_')[1]} ({pos_cnt}/{nature_cnt}/{neg_cnt})")
    plt.xlabel("steps")
    plt.plot(steps, accs, label="accuracy")
    plt.legend()
    plt.savefig(f"{log_dir}/{filename.strip('.log')}_Acc.png")

    plt.cla()
    plt.title(f"{filename.split('_')[1]} ({pos_cnt}/{nature_cnt}/{neg_cnt})")
    plt.xlabel("steps")
    plt.plot(steps, precs, label="precision")
    plt.legend()
    plt.savefig(f"{log_dir}/{filename.strip('.log')}_prec.png")

    plt.cla()
    plt.title(f"{filename.split('_')[1]} ({pos_cnt}/{nature_cnt}/{neg_cnt})")
    plt.xlabel("steps")
    plt.plot(steps, recs, label="recall")
    plt.legend()
    plt.savefig(f"{log_dir}/{filename.strip('.log')}_recall.png")

    plt.cla()
    plt.title(f"{filename.split('_')[1]} ({pos_cnt}/{nature_cnt}/{neg_cnt})")
    plt.xlabel("steps")
    plt.plot(steps, f1s, label="F1")
    plt.legend()
    plt.savefig(f"{log_dir}/{filename.strip('.log')}_f1.png")
